#!/usr/bin/python
"""\
 Converts a periodic bgf file (aka xtl file) to a seqquest input file

Copyright (c) 2003 Richard P. Muller (rmuller@sandia.gov). All rights
reserved. See the LICENSE file for licensing details.
"""

import re,string,sys
from math import sin,cos,sqrt

ang2bohr = 1.89        # Angstrom to bohr conversion
degrad = 0.01745329252 # degree to radian conversion

def get_unique_syms(geometry):
    unique_syms = []
    for (sym,x,y,z) in geometry:
        if sym not in unique_syms: unique_syms.append(sym)
    return unique_syms

def write_seqquest(lattice,geometry,filename):
    print "filename...",filename
    file = open(filename,'w')
    file.write("do setup\n")
    file.write("do iters\n")
    file.write("no force\n")  #Warning: stupid default
    file.write("no relax\n")  #Warning: stupid default
    file.write("input data\n")
    file.write("title\n")
    file.write(" SeqQuest input file generated by xtl2quest.py\n")
    file.write("dimension\n")
    file.write(" 3\n")  #Warning: stupid default
    file.write("primitive lattice vectors\n")
    a,b,c = lattice
    file.write("%15.9f %15.9f %15.9f\n" % a)
    file.write("%15.9f %15.9f %15.9f\n" % b)
    file.write("%15.9f %15.9f %15.9f\n" % c)
    file.write("grid dimensions\n")
    file.write(" 25 25 25\n")  #Warning: stupid default
    file.write("atom types\n")
    unique_syms = get_unique_syms(geometry)
    file.write("%d \n" % len(unique_syms))
    for sym in unique_syms:
        file.write("atom file\n")
        file.write("%s.atm\n" % sym)
    file.write("number of atoms in unit cell\n")
    file.write("%5d\n" % len(geometry))
    file.write("atom, type, position vector\n")
    n = 1
    for (sym,x,y,z) in geometry:
        file.write("%d %d %15.9f %15.9f %15.9f\n" %\
                   (n,unique_syms.index(sym)+1,
                    x*ang2bohr,y*ang2bohr,z*ang2bohr))
        n = n+1
    file.write("kgrid\n")
    file.write(" 4 4 4\n")
    file.write("end setup phase data\n")
    file.write("run phase data\n")
    file.write("convergence criterion\n")
    file.write(" 0.000500\n")
    file.write("end of run phase data\n")
    return

def convert_lattice(a,b,c,alpha,beta,gamma):
    """Convert to cartesian lattice vectors.
       Taken from the routine 
       /biodesign/v330/common/code/source/xtlgraf_batch/celori.f"""

    s1 = sin(alpha*degrad)
    s2 = sin(beta*degrad)
    s3 = sin(gamma*degrad)
    c1 = cos(alpha*degrad)
    c2 = cos(beta*degrad)
    c3 = cos(gamma*degrad)
    c3bar = (-c1*c2 + c3)/(s1*s2)
    sqrtarg = 1.0 - c3bar*c3bar
    if (sqrtarg <= 0):
        print "Negative argument to SQRT"
        print sqrtarg, alpha, beta, gamma
        sqrtarg = 0.0

    s3bar = sqrt(sqrtarg)

    # The general transformation from scaled to XYZ coordinates is now:
    # x = or1 * xabc
    # y = or2 * xabc + or3 * yabc
    # z = or4 * xabc + or5 * yabc + or6 * zabc

    or1 = a*s2*s3bar
    or2 = a*s2*c3bar
    or3 = b*s1
    or4 = a*c2
    or5 = b*c1
    or6 = c
    volume = a*b*c*s1*s2*s3bar

    # Compute Cartesian vectors for a, b, and c
    va = ang2bohr*or1, ang2bohr*or2, ang2bohr*or4
    vb = 0,ang2bohr*or3,ang2bohr*or5
    vc = 0,0,ang2bohr*or6

    return va,vb,vc

def cleansym(s): return string.lower(re.split('[^a-zA-Z]',s)[0])

def parse_lat_line(line):
    words = string.split(line)
    a,b,c = float(words[1]),float(words[2]),float(words[3])
    alpha,beta,gamma = float(words[4]),float(words[5]),float(words[6])
    return a,b,c,alpha,beta,gamma

def parse_at_line(line):
    words = string.split(line)
    sym = cleansym(words[2])
    x,y,z = float(words[6]),float(words[7]),float(words[8])
    return sym,x,y,z    

def read_xtl(filename):
    latpat = re.compile('CRYSTX')
    atpat = re.compile('HETATM')
    file = open(filename)
    geo = []
    while 1:
        line = file.readline()
        if not line: break
        if latpat.search(line):
            a,b,c,alpha,beta,gamma = parse_lat_line(line)
        elif atpat.search(line):
            sym,x,y,z = parse_at_line(line)
            geo.append((sym,x,y,z))
    file.close()
    lattice = convert_lattice(a,b,c,alpha,beta,gamma)
    return lattice,geo        

def xtl2quest():
    if len(sys.argv) < 1:
        print __doc__
        sys.exit()
    filename = sys.argv[1]
    lattice,geometry = read_xtl(filename)
    outfilename = string.replace(filename,'.xtl','.in')
    outfilename = string.replace(outfilename,'.bgf','.in')
    write_seqquest(lattice,geometry,outfilename)
    return

if __name__ == '__main__': xtl2quest()
